package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

// 请求记录：存储请求时间和相关信息，用于匹配响应
type RequestRecord struct {
	Timestamp time.Time // 请求时间
	SrcMAC    string    // 源 MAC
	SrcIP     string    // 源 IP
	SrcPort   uint16    // 源端口
	DstIP     string    // 目标 IP
	DstPort   uint16    // 目标端口
	Method    string    // HTTP 方法
	URL       string    // 请求 URL
	Headers   string    // 请求头
	Body      string    // 请求体
}

// 全局 map：用于存储请求记录，key 是 TCP 四元组（源IP:源端口-目标IP:目标端口）
var requestMap = make(map[string]RequestRecord)

func main() {
	// 1. 打开网络设备
	iface := os.Args[1]
	handle, err := pcap.OpenLive(iface, 1600, true, pcap.BlockForever)
	if err != nil {
		log.Fatalf("无法打开设备 %s: %v", iface, err)
	}
	defer handle.Close()

	// 2. 设置 BPF 过滤器，只捕获 HTTP 常用端口的 TCP 流量
	// if err := handle.SetBPFFilter("tcp port 80 or tcp port 8080"); err != nil {
	if err := handle.SetBPFFilter("tcp"); err != nil {
		log.Fatalf("过滤器设置失败: %v", err)
	}

	// 3. 循环捕获并解析数据包
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	for packet := range packetSource.Packets() {
		processPacket(packet)
	}
}

// processPacket 处理单个数据包，提取 HTTP 信息并匹配请求响应
func processPacket(packet gopacket.Packet) {
	// 解析链路层
	ethLayer := packet.Layer(layers.LayerTypeEthernet)
	if ethLayer == nil {
		return
	}
	eth := ethLayer.(*layers.Ethernet)

	// 解析网络层 (IPv4)
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		return
	}
	ip := ipLayer.(*layers.IPv4)

	// 解析传输层 (TCP)
	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	if tcpLayer == nil {
		return
	}
	tcp := tcpLayer.(*layers.TCP)

	// 解析应用层数据
	appLayer := packet.ApplicationLayer()
	if appLayer == nil {
		return
	}
	payload := appLayer.Payload()
	if len(payload) == 0 {
		return
	}

	// 生成 TCP 四元组（用于匹配请求响应）
	// 注意：响应的四元组是请求的反转（源IP:源端口 <-> 目标IP:目标端口）
	//requestKey := fmt.Sprintf("%s:%d-%s:%d", ip.SrcIP, tcp.SrcPort, ip.DstIP, tcp.DstPort)
	responseKey := fmt.Sprintf("%s:%d-%s:%d", ip.DstIP, tcp.DstPort, ip.SrcIP, tcp.SrcPort)

	// 区分 HTTP 请求和响应
	if isHTTPRequest(payload) {
		// 解析 HTTP 请求并存储到 map
		parseHTTPRequest(payload, eth, ip, tcp, packet.Metadata().Timestamp)
	} else if isHTTPResponse(payload) {
		// 解析 HTTP 响应并匹配对应的请求
		parseHTTPResponse(payload, eth, ip, tcp, packet.Metadata().Timestamp, responseKey)
	}
}

// isHTTPRequest 判断是否为 HTTP 请求（通过请求方法开头）
func isHTTPRequest(payload []byte) bool {
	methods := []string{"GET ", "POST ", "PUT ", "DELETE ", "HEAD ", "OPTIONS "}
	for _, m := range methods {
		if strings.HasPrefix(string(payload), m) {
			return true
		}
	}
	return false
}

// isHTTPResponse 判断是否为 HTTP 响应（通过 HTTP/ 开头）
func isHTTPResponse(payload []byte) bool {
	return strings.HasPrefix(string(payload), "HTTP/")
}

// parseHTTPRequest 解析 HTTP 请求并存储到全局 map
func parseHTTPRequest(payload []byte, eth *layers.Ethernet, ip *layers.IPv4, tcp *layers.TCP, timestamp time.Time) {
	// 分割请求头和请求体
	parts := strings.SplitN(string(payload), "\r\n\r\n", 2)
	headersStr := parts[0]
	body := ""
	if len(parts) > 1 {
		body = parts[1]
	}

	// 解析请求行
	lines := strings.Split(headersStr, "\r\n")
	if len(lines) == 0 {
		return
	}
	requestLine := lines[0]
	requestParts := strings.Split(requestLine, " ")
	if len(requestParts) < 3 {
		return
	}
	method := requestParts[0]
	path := requestParts[1]

	// 解析 Host 头
	host := ""
	scanner := bufio.NewScanner(strings.NewReader(headersStr))
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "Host: ") {
			host = strings.TrimPrefix(line, "Host: ")
			break
		}
	}
	url := fmt.Sprintf("http://%s%s", host, path)

	// 生成请求的唯一标识（TCP 四元组）
	requestKey := fmt.Sprintf("%s:%d-%s:%d", ip.SrcIP, tcp.SrcPort, ip.DstIP, tcp.DstPort)

	// 存储请求记录
	requestMap[requestKey] = RequestRecord{
		Timestamp: timestamp,
		SrcMAC:    eth.SrcMAC.String(),
		SrcIP:     ip.SrcIP.String(),
		SrcPort:   uint16(tcp.SrcPort),
		DstIP:     ip.DstIP.String(),
		DstPort:   uint16(tcp.DstPort),
		Method:    method,
		URL:       url,
		Headers:   headersStr,
		Body:      body,
	}

	// 打印请求信息（包含请求时间）
	fmt.Printf("\033[32m[HTTP 请求]\033[0m\n")
	fmt.Printf("  请求时间: %s\n", timestamp.Format("2006-01-02 15:04:05.999"))
	fmt.Printf("  源 MAC: %s\n", eth.SrcMAC)
	fmt.Printf("  源 IP: %s:%d\n", ip.SrcIP, tcp.SrcPort)
	fmt.Printf("  目标 IP: %s:%d\n", ip.DstIP, tcp.DstPort)
	fmt.Printf("  方法: %s\n", method)
	fmt.Printf("  URL: %s\n", url)
	fmt.Printf("  请求头:\n%s\n", headersStr)
	if body != "" {
		fmt.Printf("  请求体:\n%s\n", body)
	}
	fmt.Printf("-------------------------\n")
}

// parseHTTPResponse 解析 HTTP 响应并匹配对应的请求
func parseHTTPResponse(payload []byte, eth *layers.Ethernet, ip *layers.IPv4, tcp *layers.TCP, timestamp time.Time, responseKey string) {
	// 从 map 中获取对应的请求记录
	reqRecord, exists := requestMap[responseKey]
	if !exists {
		// 未找到匹配的请求（可能是请求包未被捕获，或已超时删除）
		return
	}

	// 分割响应头和响应体
	parts := strings.SplitN(string(payload), "\r\n\r\n", 2)
	headersStr := parts[0]
	body := ""
	if len(parts) > 1 {
		body = parts[1]
	}

	// 解析状态行
	lines := strings.Split(headersStr, "\r\n")
	if len(lines) == 0 {
		return
	}
	statusLine := lines[0]

	// 计算时间差
	timeDiff := timestamp.Sub(reqRecord.Timestamp)

	// 打印响应信息（包含返回时间和时间差）
	fmt.Printf("\033[34m[HTTP 响应]\033[0m\n")
	fmt.Printf("  返回时间: %s\n", timestamp.Format("2006-01-02 15:04:05.999"))
	fmt.Printf("  对应请求时间: %s\n", reqRecord.Timestamp.Format("2006-01-02 15:04:05.999"))
	fmt.Printf("  响应耗时: %s\n", timeDiff)
	fmt.Printf("  源 MAC: %s\n", eth.SrcMAC)
	fmt.Printf("  源 IP: %s:%d\n", ip.SrcIP, tcp.SrcPort)
	fmt.Printf("  目标 IP: %s:%d\n", ip.DstIP, tcp.DstPort)
	fmt.Printf("  状态行: %s\n", statusLine)
	fmt.Printf("  响应头:\n%s\n", headersStr)
	if body != "" {
		fmt.Printf("  响应体:\n%s\n", body)
	}
	fmt.Printf("=========================\n")

	// 从 map 中删除已匹配的请求（避免重复匹配）
	delete(requestMap, responseKey)
}
